# yoloe 尝试使用
# 目的 区分 各类衣服的品种
# 方式 学习使用和训练方式
import os
os.environ['CUDA_VISIBLE_DEVICES']='0'
from ultralytics import YOLOE


# 自动化流程
# 左边输入关键词
# 右侧上传图片文件夹+寻找 => 筛选

# Initialize a YOLOE model
model = YOLOE("yoloe-11l-seg.pt").cuda()  # or select yoloe-11s/m-seg.pt for different sizes
names = ["person",'garment']
model.set_classes(names, model.get_text_pe(names))


import pdb
pdb.set_trace()
from os.path import join as osj
from tqdm import tqdm
# for filename in tqdm(filenames):

dir = r'/mnt/nas/datasets/diction/ZipArchive0729'
save_clothing_dir = r'/mnt/nas/datasets/diction/ZipArchive0729_clo'
save_human_dir = r'/mnt/nas/datasets/diction/ZipArchive0729_human'
import os,shutil
filenames = os.listdir(dir)
if os.path.exists(save_clothing_dir):shutil.rmtree(save_clothing_dir)
os.makedirs(save_clothing_dir)
if os.path.exists(save_human_dir):shutil.rmtree(save_human_dir)
os.makedirs(save_human_dir)

count = 0
for entry in os.scandir( dir ):
    if entry.is_file() and entry.name.endswith('.jpg'):
        filename = entry.name
    else:continue

    count += 1
    print('\rprocess num : ',count,end='',flush=True)    
    
    filepath = os.path.join( dir , filename)


    results = model.predict(filepath)

    # 遍历 cls 和 conf
    # 优先 person 必须 conf > 0.9
    # 其次 garment 必须 conf > 0.9
    # 最后 else 

    if results[0].boxes.conf.__len__()==0:
        continue
    
    boxes = results[0].boxes
    find_person = False
    for cls,conf in zip(boxes.cls,boxes.conf):
        # 优先考虑person
        if cls == 0 and conf > 0.9:
            # move person
            shutil.copy2( osj(dir , filename) , osj(save_human_dir,filename) )
            find_person = True
            break
    if find_person:continue  # 找到person直接跳转
    for cls,conf in zip(boxes.cls,boxes.conf):
        if cls == 1:
            # move garment
            shutil.copy2( osj(dir , filename) , osj(save_clothing_dir,filename) )
            break

    # results[0].show()
    # pdb.set_trace()
    # if results[0].boxes.conf[0] < 0.9:

    #     print('未检测到person')

    #     # copy to clothing
    #     shutil.copy2( osj(dir , filename) , osj(save_clothing_dir,filename) )
    #     continue

    # Show results
    # results[0].show()
    # shutil.copy2( osj( dir,filename ) , osj( save_human_dir , filename ) )
    # pdb.set_trace()
